import asyncio
import logging
import os
import time
from typing import Optional

import ffmpeg
import numpy as np
from faster_whisper import WhisperModel

from ..config import settings
# <--- 修改: 导入 LLMService，因为 correct_text 方法需要它 --->
from ..services.llm_service import LLMService

logger = logging.getLogger(__name__)


class STTService:
    def __init__(self):
        # 1. 初始化 Whisper 模型
        model_size = "medium"
        self.model = WhisperModel(
            model_size,
            device=settings.WHISPER_DEVICE,
            compute_type="float16" if settings.WHISPER_DEVICE == "cuda" else "int8",
            download_root=os.path.join(os.path.expanduser("~"), ".cache", "whisper")
        )
        logger.info(f"Whisper {model_size} model loaded successfully")

        # <--- 修改: 在 STTService 初始化时，也创建 LLMService 实例 --->
        # 2. 初始化 LLM 服务
        self.llm_service = LLMService()
        logger.info("LLMService initialized within STTService")


    async def transcribe_audio(self, audio_data: bytes) -> str:
        """
        进行语音识别，并在得到结果后直接调用内部的校对方法。
        """
        if self.model is None:
            logger.error("Cannot transcribe audio: model not loaded")
            return "语音识别服务不可用"

        logger.info(f"Received audio data of length: {len(audio_data)} bytes")
        if len(audio_data) < 100:
            logger.warning("Received audio data is too small, likely empty or corrupted.")
            return ""

        start_time = time.time()

        try:
            process = (
                ffmpeg
                .input('pipe:', format='webm')
                .output('pipe:', format='s16le', ac=1, ar=16000)
                .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True)
            )
            pcm_data, err = process.communicate(input=audio_data)
            process.wait()

            audio_array = np.frombuffer(pcm_data, dtype=np.int16).astype(np.float32) / 32768.0

            def run_transcription():
                prompt = "以下是普通话的句子。你好，很高兴认识你。这是一个测试。希望能够正确识别。简体中文。"
                segments, _ = self.model.transcribe(
                    audio_array,
                    language="zh",
                    initial_prompt=prompt,
                    beam_size=10,
                    vad_filter=True,
                    vad_parameters=dict(min_silence_duration_ms=500)
                )
                return " ".join([segment.text for segment in segments]).strip()

            raw_text = await asyncio.to_thread(run_transcription)

            if raw_text:
                corrected_text = await self.correct_text(raw_text)
                text = corrected_text
            else:
                text = raw_text

            processing_time = time.time() - start_time
            logger.debug(f"Transcription and correction completed in {processing_time:.2f}s")

            return text
        except Exception as e:
            logger.exception(f"Speech recognition error: {e}")
            return ""

    # <--- 修改: 将 correct_text 方法直接添加到 STTService 类中 --->
    async def correct_text(self, raw_text: str, context: Optional[str] = None) -> str:
        """
        调用LLM对文本进行校对和重写。
        此方法现在是 STTService 的一部分。
        """
        system_prompt = """你是一个文本校对专家。你的任务是将用户输入的语音识别文本进行校对，修正其中的错别字、语法错误和标点符号，使其变得流畅自然。

        请严格遵守以下规则：
        1. **保持原意**：绝对不要改变句子原本想要表达的意思。
        2. **只做最小必要修改**：只修改错误的字词和语法，不要添加额外内容或进行意译。
        3. **输出要求**：只返回修正后的文本，不要有任何额外的解释、说明或标记。

        如果文本已经是流畅正确的中文，请直接返回原文本。"""

        user_input = f"请校对以下文本：{raw_text}"
        if context:
            user_input = f"上下文：{context}\n\n请根据以上上下文校对以下文本：{raw_text}"

        try:
            corrected_text = ""
            # 注意这里使用 self.llm_service
            async for chunk in self.llm_service.generate_response(
                    messages=[{"role": "user", "content": user_input}],
                    character_prompt=system_prompt,
                    stream=False
            ):
                corrected_text += chunk

            return corrected_text.strip() if corrected_text.strip() else raw_text
        except Exception as e:
            logger.error(f"LLM correction error: {e}")
            return raw_text

# 创建单例实例
stt_service = STTService()